{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Where Filtering\n",
    "This notebook demonstrates how to use where filtering to filter the data returned from get or query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromadb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using embedded DuckDB without persistence: data will be transient\n"
     ]
    }
   ],
   "source": [
    "client = chromadb.Client()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No embedding_function provided, using default embedding function: SentenceTransformerEmbeddingFunction\n"
     ]
    }
   ],
   "source": [
    "# Create a new chroma collection\n",
    "collection_name = \"filter_example_collection\"\n",
    "collection = client.create_collection(name=collection_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add some data to the collection\n",
    "collection.add(\n",
    "    embeddings=[\n",
    "        [1.1, 2.3, 3.2],\n",
    "        [4.5, 6.9, 4.4],\n",
    "        [1.1, 2.3, 3.2],\n",
    "        [4.5, 6.9, 4.4],\n",
    "        [1.1, 2.3, 3.2],\n",
    "        [4.5, 6.9, 4.4],\n",
    "        [1.1, 2.3, 3.2],\n",
    "        [4.5, 6.9, 4.4],\n",
    "    ],\n",
    "    metadatas=[\n",
    "        {\"status\": \"read\"},\n",
    "        {\"status\": \"unread\"},\n",
    "        {\"status\": \"read\"},\n",
    "        {\"status\": \"unread\"},\n",
    "        {\"status\": \"read\"},\n",
    "        {\"status\": \"unread\"},\n",
    "        {\"status\": \"read\"},\n",
    "        {\"status\": \"unread\"},\n",
    "    ],\n",
    "    documents=[\"A document that discusses domestic policy\", \"A document that discusses international affairs\", \"A document that discusses kittens\", \"A document that discusses dogs\", \"A document that discusses chocolate\", \"A document that is sixth that discusses government\", \"A document that discusses international affairs\", \"A document that discusses global affairs\"],\n",
    "    ids=[\"id1\", \"id2\", \"id3\", \"id4\", \"id5\", \"id6\", \"id7\", \"id8\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': ['id7'],\n",
       " 'embeddings': None,\n",
       " 'documents': ['A document that discusses international affairs'],\n",
       " 'metadatas': [{'status': 'read'}]}"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get documents that are read and about affairs\n",
    "collection.get(where={\"status\": \"read\"}, where_document={\"$contains\": \"affairs\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': ['id8', 'id1'],\n",
       " 'embeddings': None,\n",
       " 'documents': ['A document that discusses global affairs',\n",
       "  'A document that discusses domestic policy'],\n",
       " 'metadatas': [{'status': 'unread'}, {'status': 'read'}]}"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get documents that are about global affairs or domestic policy\n",
    "collection.get(where_document={\"$or\": [{\"$contains\": \"global affairs\"}, {\"$contains\": \"domestic policy\"}]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': [['id7', 'id2', 'id8']],\n",
       " 'embeddings': None,\n",
       " 'documents': [['A document that discusses international affairs',\n",
       "   'A document that discusses international affairs',\n",
       "   'A document that discusses global affairs']],\n",
       " 'metadatas': [[{'status': 'read'},\n",
       "   {'status': 'unread'},\n",
       "   {'status': 'unread'}]],\n",
       " 'distances': [[16.740001678466797, 87.22000122070312, 87.22000122070312]]}"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get 5 closest vectors to [0, 0, 0] that are about affairs\n",
    "# Outputs 3 docs because collection only has 3 docs about affairs\n",
    "collection.query(query_embeddings=[[0, 0, 0]], where_document={\"$contains\": \"affairs\"}, n_results=5)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "2395417914bce3169eff793a7d01bf858f95b138000d8d354eed93ead856f5e6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
